{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2a2efbb8-4d35-4df5-b519-5953f5982383",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4ae38516-b07a-465b-a080-152d4d761be1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def load_T1T2matched_datalist(args, input_file: str):\n",
    "    with open(input_file, 'r') as f:\n",
    "        fold = json.load(f)\n",
    "    print(fold.keys())\n",
    "    training_images = fold['training'] # Should be list\n",
    "    validation_images = fold['validation']\n",
    "    return {'training': training_images,\n",
    "            'validation': validation_images}\n",
    "\n",
    "def get_T1T2_dataloaders(args, modality=\"T1\", num_workers = 4):\n",
    "  \n",
    "    datalist = load_T1T2matched_datalist(args, args[\"splits\"])\n",
    "    print(datalist[\"training\"])\n",
    "    if modality == \"T1\":\n",
    "        training_datalist   = [{\"image\":subject[\"image\"][0]} for subject in datalist['training']]  \n",
    "        validation_datalist = [{\"image\":subject[\"image\"][0]} for subject in datalist['validation']] \n",
    "    elif modality == \"T2\":\n",
    "        training_datalist   = [{\"image\":subject[\"image\"][1]} for subject in datalist['training']] \n",
    "        validation_datalist = [{\"image\":subject[\"image\"][1]} for subject in datalist['validation']] \n",
    "    elif modality == \"T1T2\":\n",
    "        training_datalist, validation_datalist = datalist['training'], datalist['validation']\n",
    "    else:\n",
    "        raise ValueError(\"Unsupported modality\")\n",
    "    return training_datalist, validation_datalist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4687928f-5272-4d06-a155-e335e18a532d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['training', 'validation'])\n"
     ]
    }
   ],
   "source": [
    "args={\n",
    "'splits':'GBR_T1T2_matched_image.json'\n",
    "}\n",
    "datalist = load_T1T2matched_datalist(args,'GBR_T1T2_matched_image.json') "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "07389f30-b46f-4baf-b55d-2aa1c4efe334",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "training: 39483, validation: 1917\n"
     ]
    }
   ],
   "source": [
    "print(f\"training: {len(datalist['training'])}, validation: {len(datalist['validation'])}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d15a4384-0900-4b89-bb79-e3a36feddfc9",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1917"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "20d66576-7126-4e70-a118-220833b50b79",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub data rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_data_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_data_rate_limit=1000000.0 (bytes/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'image': '/red/ruogu.fang/UKB/data/Brain/20252_T1_NIFTI/T1_new_unzip/1977416_20252_2_0/T1/T1_brain_to_MNI.nii.gz'},\n",
       " {'image': '/red/ruogu.fang/UKB/data/Brain/20252_T1_NIFTI/T1_new_unzip/1498922_20252_2_0/T1/T1_brain_to_MNI.nii.gz'},\n",
       " {'image': '/red/ruogu.fang/UKB/data/Brain/20252_T1_NIFTI/T1_new_unzip/5929317_20252_2_0/T1/T1_brain_to_MNI.nii.gz'},\n",
       " {'image': '/red/ruogu.fang/UKB/data/Brain/20252_T1_NIFTI/T1_new_unzip/2791783_20252_2_0/T1/T1_brain_to_MNI.nii.gz'},\n",
       " {'image': '/red/ruogu.fang/UKB/data/Brain/20252_T1_NIFTI/T1_unzip/5754348_20252_2_0/T1_brain_to_MNI.nii.gz'}]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "training_datalist, validation_datalist = get_T1T2_dataloaders(args, \"T1\")\n",
    "training_datalist[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "283db78a-5fdf-4c7c-82db-5e7c38f847f2",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub data rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_data_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_data_rate_limit=1000000.0 (bytes/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'image': '/red/ruogu.fang/UKB/data/Brain/20253_T2_NIFTI/T2_unzip/1977416_20253_2_0/T2_FLAIR/T2_FLAIR_brain_to_MNI.nii.gz'},\n",
       " {'image': '/red/ruogu.fang/UKB/data/Brain/20253_T2_NIFTI/T2_unzip/1498922_20253_2_0/T2_FLAIR/T2_FLAIR_brain_to_MNI.nii.gz'},\n",
       " {'image': '/red/ruogu.fang/UKB/data/Brain/20253_T2_NIFTI/T2_unzip/5929317_20253_2_0/T2_FLAIR/T2_FLAIR_brain_to_MNI.nii.gz'},\n",
       " {'image': '/red/ruogu.fang/UKB/data/Brain/20253_T2_NIFTI/T2_unzip/2791783_20253_2_0/T2_FLAIR/T2_FLAIR_brain_to_MNI.nii.gz'},\n",
       " {'image': '/red/ruogu.fang/UKB/data/Brain/20253_T2_NIFTI/T2_unzip/5754348_20253_2_0/T2_FLAIR/T2_FLAIR_brain_to_MNI.nii.gz'}]"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "training_datalist, validation_datalist = get_T1T2_dataloaders(args, \"T2\")\n",
    "training_datalist[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "9dd382d2-cfc4-4372-ad7f-5aa3d4c2ceb0",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub data rate exceeded.\n",
      "The Jupyter server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--ServerApp.iopub_data_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "ServerApp.iopub_data_rate_limit=1000000.0 (bytes/sec)\n",
      "ServerApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'image': ['/red/ruogu.fang/UKB/data/Brain/20252_T1_NIFTI/T1_new_unzip/1977416_20252_2_0/T1/T1_brain_to_MNI.nii.gz',\n",
       "   '/red/ruogu.fang/UKB/data/Brain/20253_T2_NIFTI/T2_unzip/1977416_20253_2_0/T2_FLAIR/T2_FLAIR_brain_to_MNI.nii.gz']},\n",
       " {'image': ['/red/ruogu.fang/UKB/data/Brain/20252_T1_NIFTI/T1_new_unzip/1498922_20252_2_0/T1/T1_brain_to_MNI.nii.gz',\n",
       "   '/red/ruogu.fang/UKB/data/Brain/20253_T2_NIFTI/T2_unzip/1498922_20253_2_0/T2_FLAIR/T2_FLAIR_brain_to_MNI.nii.gz']},\n",
       " {'image': ['/red/ruogu.fang/UKB/data/Brain/20252_T1_NIFTI/T1_new_unzip/5929317_20252_2_0/T1/T1_brain_to_MNI.nii.gz',\n",
       "   '/red/ruogu.fang/UKB/data/Brain/20253_T2_NIFTI/T2_unzip/5929317_20253_2_0/T2_FLAIR/T2_FLAIR_brain_to_MNI.nii.gz']},\n",
       " {'image': ['/red/ruogu.fang/UKB/data/Brain/20252_T1_NIFTI/T1_new_unzip/2791783_20252_2_0/T1/T1_brain_to_MNI.nii.gz',\n",
       "   '/red/ruogu.fang/UKB/data/Brain/20253_T2_NIFTI/T2_unzip/2791783_20253_2_0/T2_FLAIR/T2_FLAIR_brain_to_MNI.nii.gz']},\n",
       " {'image': ['/red/ruogu.fang/UKB/data/Brain/20252_T1_NIFTI/T1_unzip/5754348_20252_2_0/T1_brain_to_MNI.nii.gz',\n",
       "   '/red/ruogu.fang/UKB/data/Brain/20253_T2_NIFTI/T2_unzip/5754348_20253_2_0/T2_FLAIR/T2_FLAIR_brain_to_MNI.nii.gz']}]"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "training_datalist, validation_datalist = get_T1T2_dataloaders(args, \"T1T2\")\n",
    "training_datalist[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "cc9b5ae1-482c-410c-bf25-4b2ecd0da9ef",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([3, 2, 3])\n",
      "tensor([[ 1,  2,  3],\n",
      "        [ 7,  8,  9],\n",
      "        [13, 14, 15]])\n",
      "torch.Size([3, 3])\n",
      "tensor([[ 4,  5,  6],\n",
      "        [10, 11, 12],\n",
      "        [16, 17, 18]])\n",
      "torch.Size([3, 3])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "x = torch.tensor([[[1, 2, 3],\n",
    "                   [4, 5, 6]],\n",
    "                  \n",
    "                  [[7, 8, 9],\n",
    "                   [10, 11, 12]],\n",
    "                  \n",
    "                  [[13, 14, 15],\n",
    "                   [16, 17, 18]]])\n",
    "\n",
    "print(x.shape)\n",
    "selected_elements = x[:, 0]\n",
    "print(selected_elements)\n",
    "print(selected_elements.shape)\n",
    "selected_elements = x[:, 1]\n",
    "print(selected_elements)\n",
    "print(selected_elements.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c82ff1d-b6ec-4065-b585-465bdbd840bc",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "monai-0.9.1-local",
   "language": "python",
   "name": "monai-0.9.1-local"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
